package com.example.scribe.backend.handlers;

import com.example.scribe.SyntaxNode;
import com.example.scribe.backend.BytecodeOps;
import com.example.scribe.backend.ExpressionEmitter;
import com.example.scribe.backend.ExpressionNodeHandler;
import com.example.scribe.backend.LocalVariableAllocator;
import net.bytebuddy.implementation.bytecode.StackManipulation;
import net.bytebuddy.implementation.bytecode.StackManipulation.Compound;

/**
 * Handler for bitwise expressions (bitwise_and, bitwise_or, bitwise_xor, shift).
 */
public final class BitwiseExpressionHandler implements ExpressionNodeHandler {
    private final ExpressionEmitter exprEmitter;

    public BitwiseExpressionHandler(ExpressionEmitter exprEmitter) {
        this.exprEmitter = exprEmitter;
    }

    @Override
    public boolean canHandle(String nodeType) {
        return "bitwise_and".equals(nodeType) || "bitwise_or".equals(nodeType) || 
               "bitwise_xor".equals(nodeType) || "shift".equals(nodeType);
    }

    @Override
    public StackManipulation handle(SyntaxNode node, String generatedClassName, LocalVariableAllocator locals) {
        if (node == null) {
            return StackManipulation.Trivial.INSTANCE;
        }

        String nodeType = node.type();
        
        // For shift nodes, the structure is: multiplication (operator multiplication)*
        // If there's only one multiplication child, just emit it (no shift operation)
        // If there are operator and operands, perform the shift operation
        if ("shift".equals(nodeType)) {
            // If shift node has only one child, it's just a multiplication (no shift operator)
            // In this case, emit the child directly (pass through to multiplication/unary)
            if (node.children().size() == 1) {
                SyntaxNode onlyChild = node.children().get(0);
                // Recursively emit the child, which should be a multiplication or unary node
                return exprEmitter.emit(onlyChild, generatedClassName, locals);
            }
            
            // Find operator and operands in shift node
            SyntaxNode leftOperand = null;
            SyntaxNode rightOperand = null;
            String op = null;
            
            // Look for operator terminal and multiplication operands
            for (int i = 0; i < node.children().size(); i++) {
                SyntaxNode child = node.children().get(i);
                String childType = child.type();
                
                // Check for operator terminal
                if (childType.startsWith("Terminal:")) {
                    String text = child.text();
                    if ("<<".equals(text) || ">>".equals(text) || ">>>".equals(text)) {
                        op = text;
                        // Left operand should be before operator, right operand after
                        if (i > 0) {
                            leftOperand = node.children().get(i - 1);
                        }
                        if (i + 1 < node.children().size()) {
                            rightOperand = node.children().get(i + 1);
                        }
                        break;
                    }
                }
            }
            
            if (leftOperand == null || rightOperand == null || op == null) {
                // If we can't find operator and operands, return Trivial
                // This should not happen if the AST is correctly structured
                return StackManipulation.Trivial.INSTANCE;
            }
            
            StackManipulation left = exprEmitter.emit(leftOperand, generatedClassName, locals);
            StackManipulation right = exprEmitter.emit(rightOperand, generatedClassName, locals);
            
            if (left == StackManipulation.Trivial.INSTANCE || right == StackManipulation.Trivial.INSTANCE) {
                return StackManipulation.Trivial.INSTANCE;
            }
            
            if ("<<".equals(op)) {
                return new Compound(left, right, BytecodeOps.ishl());
            } else if (">>".equals(op)) {
                return new Compound(left, right, BytecodeOps.ishr());
            } else if (">>>".equals(op)) {
                return new Compound(left, right, BytecodeOps.iushr());
            }
            return StackManipulation.Trivial.INSTANCE;
        }
        
        // For bitwise_and, bitwise_or, bitwise_xor, expect at least 3 children
        if (node.children().size() < 3) {
            return StackManipulation.Trivial.INSTANCE;
        }
        
        StackManipulation left = exprEmitter.emit(node.children().get(0), generatedClassName, locals);
        StackManipulation right = exprEmitter.emit(node.children().get(2), generatedClassName, locals);

        if ("bitwise_and".equals(nodeType)) {
            // Check for & operator
            String op = null;
            if (node.children().size() >= 2) {
                SyntaxNode mid = node.children().get(1);
                op = mid.text();
            }
            if ("&".equals(op)) {
                return new Compound(left, right, BytecodeOps.iand());
            }
            return StackManipulation.Trivial.INSTANCE;
        } else if ("bitwise_or".equals(nodeType)) {
            // Check for | operator
            // bitwise_or structure: bitwise_xor ( '|' bitwise_xor )*
            // So we need to find the | operator in children
            String op = null;
            StackManipulation leftOp = null;
            StackManipulation rightOp = null;
            for (int i = 0; i < node.children().size(); i++) {
                SyntaxNode child = node.children().get(i);
                if (child.type().startsWith("Terminal:")) {
                    String text = child.text();
                    if ("|".equals(text)) {
                        op = text;
                        // Update left and right operands based on operator position
                        if (i > 0) {
                            leftOp = exprEmitter.emit(node.children().get(i - 1), generatedClassName, locals);
                        }
                        if (i + 1 < node.children().size()) {
                            rightOp = exprEmitter.emit(node.children().get(i + 1), generatedClassName, locals);
                        }
                        break;
                    }
                }
            }
            // Fallback: check middle child (original logic)
            if (op == null && node.children().size() >= 2) {
                SyntaxNode mid = node.children().get(1);
                op = mid.text();
                // Use original left and right
                leftOp = left;
                rightOp = right;
            }
            if ("|".equals(op)) {
                if (leftOp == null) leftOp = left;
                if (rightOp == null) rightOp = right;
                if (leftOp == StackManipulation.Trivial.INSTANCE || rightOp == StackManipulation.Trivial.INSTANCE) {
                    return StackManipulation.Trivial.INSTANCE;
                }
                return new Compound(leftOp, rightOp, BytecodeOps.ior());
            }
            return StackManipulation.Trivial.INSTANCE;
        } else if ("bitwise_xor".equals(nodeType)) {
            // Check for ^ operator
            String op = null;
            if (node.children().size() >= 2) {
                SyntaxNode mid = node.children().get(1);
                op = mid.text();
            }
            if ("^".equals(op)) {
                return new Compound(left, right, BytecodeOps.ixor());
            }
            return StackManipulation.Trivial.INSTANCE;
        }

        return StackManipulation.Trivial.INSTANCE;
    }
}

